"""
工具函数模块

主要功能：
1. 训练过程可视化（Tensorboard支持）
2. 多种评估指标的计算和记录
3. 模型状态字典的处理
4. 参数配置的表格化展示

评估指标包括：
- AUC (Area Under Curve)：衡量分类器的区分能力
- ACC (Accuracy)：分类准确率
- LogLoss：对数损失，反映预测的确信度
- TestLoss：测试集上的损失值
"""

from metrics.accuracy_metric import AccuracyMetric
from metrics.metric import Metric
from metrics.test_loss_metric import TestLossMetric
from metrics.auc_metric import AUCMetric
from metrics.logloss_metric import LOGLOSSMetric
import torch
from torch.utils.tensorboard import SummaryWriter
import logging

class plt_tensorboard():
    """
    Tensorboard可视化类
    
    功能：
    - 记录训练过程中的各种指标
    - 支持多种评估指标的累积和重置
    - 提供指标的可视化和报告功能
    """
    def __init__(self, args):
        """
        初始化Tensorboard写入器和评估指标
        Args:
            args: 参数配置对象
        """
        wr = SummaryWriter(log_dir=f'runs/{args.name}')
        self.tb_writer = wr

        # 初始化多个评估指标
        self.metrics = [
            AccuracyMetric(),    # 准确率
            TestLossMetric(),    # 测试损失
            AUCMetric(),         # AUC曲线下面积
            LOGLOSSMetric()      # 对数损失
        ]

    def accumulate_metrics(self, outputs, labels, loss):
        """
        累积一个批次的评估指标
        Args:
            outputs: 模型输出
            labels: 真实标签
            loss: 损失值
        """
        self.metrics[0].accumulate_on_batch([outputs, labels])  # 累积准确率
        self.metrics[1].accumulate_on_batch(loss)              # 累积测试损失
        self.metrics[2].accumulate_on_batch([outputs, labels])  # 累积AUC
        self.metrics[3].accumulate_on_batch([outputs, labels])  # 累积对数损失

    def reset_metrics(self):
        """重置所有评估指标，通常在每个epoch开始时调用"""
        for metric in self.metrics:
            metric.reset_metric()

    def report_metrics(self, step, tb_writer=None, tb_prefix='Metric/'):
        """
        报告当前的评估指标
        Args:
            step: 当前步数/轮数
            tb_writer: Tensorboard写入器
            tb_prefix: 指标前缀
        Returns:
            AUC值，用于模型选择
        """
        metric_text = []
        # 记录所有指标并写入Tensorboard
        for metric in self.metrics:
            metric_text.append(str(metric))
            metric.plot(tb_writer, step, tb_prefix=tb_prefix)
        
        try:
            # 打印主要指标
            print(f"AUC: {self.metrics[2].get_value()['value']}, "
                  f"ACC: {self.metrics[0].get_value()['value']}, "
                  f"LogLoss: {self.metrics[3].get_value()['value']}")
        except:
            print("error")
        return self.metrics[2].get_value()['value']  # 返回AUC值

def remove_prefix(state_dict, prefix):
    """
    移除模型状态字典中的前缀
    
    用途：
    - 处理多GPU训练时的模型权重加载
    - 移除'module.'前缀以适配单GPU使用
    
    Args:
        state_dict: 模型状态字典
        prefix: 要移除的前缀（通常是'module.'）
    Returns:
        处理后的状态字典
    """
    print('remove prefix \'{}\''.format(prefix))
    f = lambda x: x.split(prefix, 1)[-1] if x.startswith(prefix) else x
    return {f(key): value for key, value in state_dict.items()}

def create_table(params):
    """
    创建参数配置的Markdown表格
    
    功能：
    - 将参数配置转换为易读的表格形式
    - 用于记录实验配置和生成报告
    
    Args:
        params: 参数配置对象
    Returns:
        Markdown格式的表格字符串
    """
    data = "| name | value | \n |-----|-----|"
    params = params.__dict__
    for key, value in params.items():
        data += '\n' + f"| {key} | {value} |"

    return data